"""最小二乘法"""
import numpy as np
import matplotlib.pyplot as plt

#https://zhuanlan.zhihu.com/p/38128785/

def fun2ploy(x,n):
    '''
    数据转化为[x^0,x^1,x^2,...x^n]
    首列变1
    '''
    lens = len(x)
    X = np.ones([1,lens])
    for i in range(1,n):
        X = np.vstack((X,np.power(x,i)))#按行堆叠
    return X  



def leastseq_byploy(x,y,ploy_dim):
    '''
    最小二乘求解
    '''
    X = fun2ploy(x,ploy_dim) #m*n矩阵
    #直接求解
    Xt = X.transpose();#转置变成列向量 X=[[x^0][x^1][x^2][x^3]]
    XXt=X.dot(Xt);#矩阵乘,X*X^T
    XXtInv = np.linalg.inv(XXt)#求逆 (X*X^T)^-1
    XXtInvX = XXtInv.dot(X) #(X*X^T)^-1*X
    coef = XXtInvX.dot(y.T)
    y_est = Xt.dot(coef)
    return y_est,coef

def fit_fun(x):  ## 如 p=numpy.poly1d([1,2,3])  生成  $1x^2+2x^1+3x^0$*,coef
    return np.sin(x)

if __name__ == '__main__':
    data_num = 100 # m=100,(x,y)有100对
    ploy_dim =10  #拟合参数个数，即权重数量,y= A1+A2*X+A3*X^2+A4*X^3+...+A10*X^9
    noise_scale = 0.2
    ## 数据准备
    x = np.array(np.linspace(-2*np.pi,2*np.pi,data_num))   #数据 
    y = fit_fun(x)+noise_scale*np.random.rand(1,data_num)  #添加噪声np.random.rand(1,data_num) ,1*data_num矩阵

    # 最小二乘拟合
    [y_est,coef] = leastseq_byploy(x,y,ploy_dim)
    
    #显示拟合结果
    org_data = plt.scatter(x,y,color="r",marker='o',s = 50)
    est_data = plt.scatter(x,y_est,color="g",linewidth= 2)
    
    plt.xlabel("X")
    plt.ylabel("Y")
    plt.title("Fit funtion with leastseq method")
    plt.legend(["Noise data","Fited function"])
    plt.show()